##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(ggpubr)
library(ggsci)
library(patchwork)
library("scales")

##########################################################################################

option_list <- list(
    make_option(c("--gene"), type = "character") ,
    make_option(c("--input_file"), type = "character") ,
    make_option(c("--ccf_file"), type = "character") ,
    make_option(c("--maf_public_file"), type = "character") ,
    make_option(c("--sample_public_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    input_file <- paste(work_dir,"/baseTable/STAD_Info.addBurden.MSI_MSS.tsv",sep="")
    ccf_file <- paste(work_dir,"/mutationTime/result/All_CCF_mutTime.tsv",sep="")
    images_path <- paste(work_dir,"/images/mutBurden",sep="")

    maf_public_file <- paste(work_dir,"/maf_public/All_use.addVAF.maf",sep="")
    sample_public_file <- paste(work_dir,"/public_ref/combine/MutationInfo.combine.tsv",sep="")
}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

gene <- opt$gene
input_file <- opt$input_file
ccf_file <- opt$ccf_file
images_path <- opt$images_path
maf_public_file <- opt$maf_public_file
sample_public_file <- opt$sample_public_file

dir.create(images_path , recursive = T)

###########################################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")
col <- col[1:3]

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
from_order <- c("NJMU" , "EastAsian" , "TCGA")

###########################################################################################

dat_sample <- data.frame(fread( input_file ))
dat_ccf <- data.frame(fread( ccf_file ))

dat_sample_public <- data.frame(fread( sample_public_file ))
dat_maf_public <- data.frame(fread( maf_public_file ))
dat_sample_public <- subset( dat_sample_public , From != "NJMU" )

###########################################################################################
## 看reucrrent突变
dat_ccf <- subset( dat_ccf , Hugo_Symbol==gene & Variant_Classification %in% Variant_Types )
dat_ccf$vid <- paste(dat_ccf$Chr , dat_ccf$Start_Position , dat_ccf$REF , dat_ccf$ALT , sep = ":")
hotspot <- which(table( dat_ccf$vid) > 5)
dat_ccf_recurrent <- subset( dat_ccf , vid %in% names(hotspot) )
dat_ccf_nonrecurrent <- subset( dat_ccf , !(vid %in% names(hotspot) ))

dat_maf_public <- subset( dat_maf_public , Hugo_Symbol==gene & Variant_Classification %in% Variant_Types )
dat_maf_public$vid <- paste(dat_maf_public$Chromosome , dat_maf_public$Start_position , dat_maf_public$Reference_Allele , dat_maf_public$Tumor_Seq_Allele2 , sep = ":")
dat_maf_public_recurrent <- subset( dat_maf_public , vid %in% names(hotspot) )
dat_maf_public_nonrecurrent <- subset( dat_maf_public , !(vid %in% names(hotspot)) )

dat_sample_use <- dat_sample
dat_sample_public_use <- dat_sample_public

###########################################################################################

for( type in c("reccurent" , "nonreccurent") ){

    if(type == "reccurent"){
        dat_ccf <- dat_ccf_recurrent
        dat_maf_public <- dat_maf_public_recurrent
    }else if(type == "nonreccurent"){
        dat_ccf <- dat_ccf_nonrecurrent
        dat_maf_public <- dat_maf_public_nonrecurrent
    }

    dat_sample <- dat_sample_use
    dat_sample_public <- dat_sample_public_use

    dat_sample$DriverMut <- ifelse( dat_sample$Tumor %in% dat_ccf$Sample , "Mut" , "NoMut"  )
    dat_sample_public$DriverMut <- ifelse( dat_sample_public$Tumor %in% dat_maf_public$Tumor_Sample_Barcode , "Mut" , "NoMut"  )

    ###########################################################################################
    ## 去除MSI
    dat_sample <- subset( dat_sample , MS_Type != "MSI" | Class == "IM" )
    dat_sample_public <- subset( dat_sample_public , MS_Type != "MSI")

    ###########################################################################################
    ## 多个样本负荷取中位数
    dat_plot2 <- dat_sample %>%
    group_by( Patient , Class , DriverMut ) %>%
    summarize( BurdenExon = median(BurdenExon) )
    dat_plot2$From <- "NJMU"

    dat_sample_public <- dat_sample_public[,c("Tumor" , "Class" , "DriverMut" , "BurdenExon" , "From")]
    colnames(dat_sample_public)[1] <- "Patient"

    dat_plot2 <- rbind( dat_plot2 , dat_sample_public )
    dat_plot2$Class <- factor( dat_plot2$Class , levels = names(col) , order = T )
    dat_plot2$From <- factor( dat_plot2$From , levels = from_order , order = T )

    ###########################################################################################

    trans <- function(num){
        up <- floor(log10(num))
        down <- round(num*10^(-up),2)
        text <- paste("p == ",down," %*% 10","^",up)
        return(text)
    }

    #### 总的
    dat_plot_tmp <- dat_plot2
    dat_plot_tmp$MutBurden_use <- dat_plot_tmp$BurdenExon
    dat_plot_tmp$p_text <- ""
    for( class in unique(dat_plot_tmp$Class) ){

        tmp <- data.frame(subset(dat_plot_tmp , Class == class))

        if(nrow(tmp[which(tmp$DriverMut=="Mut"),]) > 0){
            p <-  wilcox.test( as.numeric(tmp[which(tmp$DriverMut=="Mut"),"MutBurden_use"]) , as.numeric(tmp[which(tmp$DriverMut=="NoMut"),"MutBurden_use"] ))$p.value
            if( p < 0.001 ){
                p_text <- trans(p)
            }else{
                p_text <- paste0( "p == " , round(as.numeric(p) , 3) ) 
            }

            dat_plot_tmp[which(dat_plot_tmp$Class == class), "p_text"] <- p_text
        }
    }

    plot <- ggplot( dat_plot_tmp , aes( x = DriverMut , y = MutBurden_use , color = DriverMut ) ) +
        geom_boxplot(alpha =1 , outlier.color=NA , size = 1.5 , width = 0.6) +
        geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
        facet_grid(.~Class,space='free_x',scales='free_x') +
        scale_color_npg() +
        xlab(NULL) +
        ylab("Mutation rate per MB")+
        theme_bw() +
        scale_y_continuous(
                limits = c(0,25) ,
                breaks = c(1 , 2 , 3 , 4 , seq(0,25,5)),
                trans = sqrt_trans()
                ) +
        geom_text(aes(label=p_text , y = 25 ,x = 1.5),parse = TRUE,size=4 , color = "black") +
        theme(panel.background = element_blank(),#设置背影为白色#清除网格线
            legend.position ='none',
            legend.title = element_blank() ,
            panel.grid.major=element_line(colour=NA),
            plot.title = element_text(size = 8,color="black",face='bold'),
            legend.text = element_text(size = 10,color="black",face='bold'),
            axis.text.y = element_text(size = 10,color="black",face='bold'),
            axis.title.x = element_text(size = 10,color="black",face='bold'),
            axis.title.y = element_text(size = 10,color="black",face='bold'),
            axis.ticks.x = element_blank(),
            axis.text.x = element_text(size = 10,color="black",face='bold') ,
            axis.line = element_line(size = 0.5)) 

    out_name <- paste0( images_path , "/mutBurden.",gene,".cds.",type,".pdf" )  
    ggsave(file=out_name,plot=plot,width=6,height=5)

    ###########################################################################################
    ## 不同来源
    dat_plot_tmp <- subset(dat_plot2 , Class!="IM")
    dat_plot_tmp$MutBurden_use <- dat_plot_tmp$BurdenExon
    dat_plot_tmp$p_text <- ""
    for( from in unique(dat_plot_tmp$From) ){
        for( class in unique(dat_plot_tmp$Class) ){
            tmp <- data.frame(subset(dat_plot_tmp , Class == class & From == from))

            if(nrow(tmp[which(tmp$DriverMut=="Mut"),]) > 0){
                p <-  wilcox.test( as.numeric(tmp[which(tmp$DriverMut=="Mut"),"MutBurden_use"]) , as.numeric(tmp[which(tmp$DriverMut=="NoMut"),"MutBurden_use"] ))$p.value

                if( p < 0.001 ){
                    p_text <- trans(p)
                }else{
                    p_text <- paste0( "p == " , round(as.numeric(p) , 3) ) 
                }

                dat_plot_tmp[which(dat_plot_tmp$Class == class & dat_plot_tmp$From == from), "p_text"] <- p_text
            }
        }
    }

    plotFunction <- function(dat_plot_tmp_from = dat_plot_tmp_from , from = from){
        plot <- ggplot( dat_plot_tmp_from , aes( x = DriverMut , y = MutBurden_use , color = DriverMut ) ) +
            geom_boxplot(alpha =1 , outlier.color=NA , size = 1.5 , width = 0.6) +
            geom_jitter(position = position_jitter(0.17) , size = 1 , alpha = 0.7) +
            facet_grid(.~Class,space='free_x',scales='free_x') +
            scale_color_npg() +
            xlab(NULL) +
            scale_y_continuous(
                limits = c(0,25) ,
                breaks = c(1 , 2 , 3 , 4 , seq(0,25,5)),
                trans = sqrt_trans()
                ) +
            ylab("Mutation rate per MB")+
            theme_bw() +
            labs(title=from) +
            geom_text(aes(label=p_text , y = 25 ,x = 1.5),parse = TRUE,size=3 , color = "black") +
            theme(panel.background = element_blank(),#设置背影为白色#清除网格线
                legend.position ='none',
                legend.title = element_blank() ,
                panel.grid.major=element_line(colour=NA),
                plot.title = element_text(size = 12,color="black",face='bold',hjust=0.5,vjust=0.5),
                legend.text = element_text(size = 10,color="black",face='bold'),
                axis.text.y = element_text(size = 10,color="black",face='bold'),
                axis.title.x = element_text(size = 10,color="black",face='bold'),
                axis.title.y = element_text(size = 10,color="black",face='bold'),
                axis.ticks.x = element_blank(),
                axis.text.x = element_text(size = 10,color="black",face='bold') ,
                axis.line = element_line(size = 0.5)) 
        
        return(plot)
    }

    from <- "NJMU"
    dat_plot_tmp_from <- subset( dat_plot_tmp , From == from )
    p1 <- plotFunction(dat_plot_tmp_from = dat_plot_tmp_from , from = from)

    from <- "EastAsian"
    dat_plot_tmp_from <- subset( dat_plot_tmp , From == from )
    p2 <- plotFunction(dat_plot_tmp_from = dat_plot_tmp_from , from = from)

    from <- "TCGA"
    dat_plot_tmp_from <- subset( dat_plot_tmp , From == from )
    p3 <- plotFunction(dat_plot_tmp_from = dat_plot_tmp_from , from = from)

    p <- p1 + p2 + p3

    out_name <- paste0( images_path , "/mutBurden.",gene,".from.cds.",type,".pdf" )  
    ggsave(file=out_name,plot=p,width=8,height=5)
}